#include "assert.h"
#include <random>
#include <algorithm>
#include <set>
#include <cmath>
#include <numeric>
#include <time.h>
#include <boost/math/distributions/normal.hpp>
#include "Sample.h"

/**
 * @brief 计算样本大小
 *
 * 根据给定的最大相对误差、置信水平和总体标准差，计算所需的样本大小。
 *
 * @param maxRelativeError 最大相对误差
 * @param confidenceLevel 置信水平
 * @param populationStdDev 总体标准差
 *
 * @return 所需的样本大小，向上取整为整数
 */
unsigned calculateSampleSize(double maxRelativeError, double confidenceLevel, double populationStdDev) {
    // 将置信水平转换为Z分数（例如，95%置信水平对应约1.96的Z分数）
    boost::math::normal_distribution<> dist(0, 1);
    double zScore = boost::math::quantile(dist, (1 + confidenceLevel) / 2);

    // 根据公式计算样本量
    return (unsigned)((zScore * populationStdDev / maxRelativeError) * (zScore * populationStdDev / maxRelativeError));
}
double calculateMaxRelativeError(unsigned sampleSize, double confidenceLevel, double populationStdDev) {
    // 将置信水平转换为Z分数（例如，95%置信水平对应约1.96的Z分数）
    boost::math::normal_distribution<> dist(0, 1);
    double zScore = boost::math::quantile(dist, (1 + confidenceLevel) / 2);

    // 计算最大相对误差
    return (zScore * populationStdDev) / std::sqrt(sampleSize);
}

namespace sample {

    std::set<unsigned> generateUniqueRandomNumbers(unsigned n, unsigned m) {
        assert(n >= m && m > 0);

        // 初始化随机数生成器
        std::mt19937 gen((unsigned)time(nullptr));
        std::uniform_int_distribution<> dis(0, n-1);

        std::set<unsigned> uniqueNumbers; // 使用set自动去重
        while(uniqueNumbers.size() < m) {
            uniqueNumbers.insert(dis(gen));
        }

        return uniqueNumbers;
    }


    bool SampleDesign::calcSampleSize() {
        if (metric_.empty()) {
            // should set the metric will be sampled.
            return false;
        }
        if (relativeError_ < 0.01) {
            return false;
        }
        
        sampleSize_ = calculateSampleSize(relativeError_, 0.95, sampleFrame_.calculatePopulationStdDev(metric_));
        if (sampleSize_ > sampleFrame_.getRowCount()) {
            sampleSize_ = sampleFrame_.getRowCount();
        }
        return true;
    }
    
    bool SampleDesign::calcRelativeError()
    {
        if (sampleSize_ == 0) {
            return false;
        }
        
        relativeError_ = calculateMaxRelativeError(sampleSize_, 0.95, sampleFrame_.calculatePopulationStdDev(metric_));
        
        return true;
    }
		
    SampleResult* RandomSample::sample() {
        if (sampleSize_ == 0 && 0 == relativeError_) {
            return nullptr;
        }
        if (metric_.empty()) {
            return nullptr;
        }

        SampleResult sampleResult;
        sampleResult.sampleDesign = this;
        
        // 随机生成n个不重复的整数
        std::set<unsigned> uniqueNumbers = generateUniqueRandomNumbers(sampleFrame_.getRowCount(), this->sampleSize_);
        sampleResult.results["random"] = uniqueNumbers;
        
        return &SampleManager::getInstance().addSampleResult(sampleResult);    
    }

    /**
     * @brief 直线等距采样
     *
     * 根据给定的总数 N、采样数量 n 和起始位置 r，进行直线等距采样。
     *
     * @param N 总数
     * @param n 采样数量
     * @param r 起始位置
     */
    void linearSystematicSampling(unsigned N, unsigned n, unsigned r, std::set<unsigned>& samples) {
        unsigned k = N / n;
        for (unsigned i = 0; i < n; ++i) {
            unsigned sample = (r + i * k);
            if (sample >= N) {
                continue;
            }
            samples.insert(sample);
        }
    }

    /**
     * @brief 循环等距采样
     *
     * 根据给定的总数 N、采样数量 n 和起始位置 r，进行循环等距采样。
     *
     * @param N 总数
     * @param n 采样数量
     * @param r 起始位置
     */
    void circularSystematicSampling(unsigned N, unsigned n, unsigned r, std::set<unsigned>& samples)
    {
        assert(N >= n);
        unsigned k = N / n;
        for (unsigned i = 0; i < n; ++i) {
            unsigned sample = (r + i * k) % N; // 循环采样确保样本编号在0到N-1之间
            samples.insert(sample);
        }
    }

    SampleResult* SystemSample::sample()
    {
        if (sampleSize_ == 0 && 0 == relativeError_) {
            return nullptr;
        }
        if (metric_.empty()) {
            return nullptr;
        }

        SampleResult sampleResult;
        sampleResult.sampleDesign = this;
        
        srand((unsigned)time(nullptr));
        unsigned r = rand() % sampleFrame_.getRowCount() + 1; ;
        std::set<unsigned> uniqueNumbers; // 使用set自动去重
        if (sortType_ == Linear) {
            linearSystematicSampling(sampleFrame_.getRowCount(), sampleSize_, r, uniqueNumbers);
            sampleResult.results["linear"] = uniqueNumbers;
        }
        else {
            circularSystematicSampling(sampleFrame_.getRowCount(), sampleSize_, r, uniqueNumbers);    
            sampleResult.results["circular"] = uniqueNumbers;
        } 
        
        return &SampleManager::getInstance().addSampleResult(sampleResult); 
    }
    
    void StratifiedSample::calcLayerPopulationSize()
    {
        if (metric_.empty()) {
            // should set the metric will be sampled.
            return;
        }

        // 确定分层，并对每层进行采样
        // 确定每层的总体数量
        layerInfos_.clear();
        if (sampleFrame_.getColumnType(layerMetric_) == Metric::TypeString) {
            auto c = sampleFrame_.getColumn(layerMetric_);
            assert(c != nullptr);
            for(unsigned i = 0; i < c->size(); ++i) {
                string& metric = std::get<string>((*c)[i]);
                string m = (layerMetricSize_ >= metric.size() || layerMetricSize_ == 0) ? metric : metric.substr(0, layerMetricSize_);
                LayerInfo& li = layerInfos_[m];
                li.populationNos.push_back(i);
            }
        }
        else {
            auto popurationCount = sampleFrame_.getRowCount();  // 总单元数
            for(unsigned i = 0; i<layerCount_ ; ++i) {
                LayerInfo& li = layerInfos_[std::to_string(i+1)];
                li.populationSize = (i== layerCount_-1) ? 
                    (popurationCount - (layerCount_ - 1) * (popurationCount / layerCount_)) :
                    (popurationCount / layerCount_);
            }                
        }
    }

    // 剩余的分配
    void calcRemainingAllocation(unsigned remainingSamples, const std::vector<unsigned>& layerSizes, std::vector<unsigned>& allocation)
    {
        for(size_t i = 0; i < layerSizes.size() && remainingSamples > 0; ++i) {
            assert(allocation[i] <= layerSizes[i]);

            if (remainingSamples + allocation[i] <= layerSizes[i]) {
                allocation[i] += remainingSamples;
                break;
            } 
            else {
                remainingSamples -= layerSizes[i] - allocation[i];
                allocation[i] = layerSizes[i];                 
            }
        }
    }

    // 平均分配中，如果有剩余的，放到下一层里面（不超过该层总体单位数）
    std::vector<unsigned> averageAllocation(const std::vector<unsigned>& layerSizes, unsigned totalSamples) {
        std::vector<unsigned> allocation(layerSizes.size(), 0);
        int remainingSamples = totalSamples;

        // 第一步：初始平均分配，但不分配超过每层的大小
        for(size_t i = 0; i < layerSizes.size(); ++i) {
            allocation[i] = std::min(remainingSamples / (layerSizes.size() - i), layerSizes[i]);
            remainingSamples -= allocation[i];
        }

        // 第二步：如果还有剩余样本，尝试向后续层分配直至无剩余或所有层达到上限
        calcRemainingAllocation(remainingSamples, layerSizes, allocation);

        return allocation;
    }

// 如果各层的总体数量已知，可以根据各层的数量按比例分配样本数。
std::vector<unsigned> proportionalAllocation(const std::vector<unsigned>& layerSizes, unsigned totalSamples) {
    double totalPopulation = 0;
    for(unsigned size : layerSizes) totalPopulation += size;
    
    std::vector<unsigned> proportionalAllocations;
    for (unsigned size : layerSizes) {
        double temp = size * totalSamples;
        proportionalAllocations.push_back((unsigned)(temp / totalPopulation));
    }

    std::vector<unsigned> allocation(layerSizes.size());
    unsigned remainingSamples = totalSamples;
    for(size_t i = 0; i < layerSizes.size(); ++i) {
        allocation[i] = proportionalAllocations[i];
        remainingSamples -= allocation[i];
    }
    
    // 分配剩余样本（可选策略）
    for(unsigned i = 0; i < remainingSamples; ++i) {
        allocation[i%allocation.size()]++;
    }
    
    return allocation;
}

// 假设每个层有一个与之相关的调查指标值（比如重要性、变异性等），并据此按比例分配样本。
void indicatorProportionalAllocation(const std::vector<double>& metricValues,
    const std::vector<unsigned>& layerSizes,
    unsigned totalSamples, 
    std::vector<unsigned>& sampleCounts) 
{
    double totalMetric = 0.0;
    for(double value : metricValues) totalMetric += value;

    for(size_t i = 0; i < metricValues.size(); ++i) {
        // 计算基于指标值的等比例样本数，同时确保不超过该层的单元数
        double tmp = totalSamples* metricValues[i];
        unsigned calculatedSample = (unsigned)(tmp / totalMetric);
        sampleCounts.push_back(std::min(calculatedSample, layerSizes[i])); // 添加约束不超过单元数
    }

    // 检查总分配样本数，确保不超过总样本数，如有必要进行调整
    unsigned sum = 0;
    for(unsigned count : sampleCounts) sum += count;
    assert(totalSamples >= sum);
    calcRemainingAllocation(totalSamples-sum, layerSizes, sampleCounts);
}

/**
 * @brief 带约束的 Neyman 分配
 *
 * 根据给定的层大小、方差和总样本数，进行 Neyman 分配，并考虑约束条件。
 *
 * @param layerSizes 层大小
 * @param variances 方差
 * @param totalSamples 总样本数
 * @param sampleCounts 分配后的样本数（输出）
 */
void neymanAllocation(const std::vector<unsigned>& layerSizes, const std::vector<double>& variances, unsigned totalSamples, std::vector<unsigned>& sampleCounts) {
    double totalVariance = 0.0;
    for(size_t i = 0; i < layerSizes.size(); ++i) {
        totalVariance += variances[i] * layerSizes[i];
    }

    for(size_t i = 0; i < layerSizes.size(); ++i) {
        unsigned calculatedSample = (unsigned) (totalSamples * sqrt(variances[i] * layerSizes[i] / totalVariance));
        // 确保分配的样本数不超过该层的实际单元数
        sampleCounts.push_back(std::min(calculatedSample, (layerSizes[i])));
    }

    // 检查总样本数是否剩余
    unsigned sum = 0;
    for(unsigned count : sampleCounts) sum += count;    
    assert(totalSamples >= sum);
    calcRemainingAllocation(totalSamples-sum, layerSizes, sampleCounts);
}

    
    void StratifiedSample::calcLayerSampleSizeByTotalSize(unsigned totalSampleSize, const string& metric, SampleConfigType type)
    {
        if (totalSampleSize <= 0)
            return;
        sampleSize_ = totalSampleSize;
        sampleType_ = type;

        vector<unsigned> layerPopulationSizes, layerSampleSizes;
        for(auto layer : layerInfos_) {
            layerPopulationSizes.push_back(isStringLayering() ? layer.second.populationNos.size() : layer.second.populationSize);
        }
        switch (type)
        {
        case SampleConfigTypeEven:            
            layerSampleSizes = averageAllocation(layerPopulationSizes, totalSampleSize);
            break;
        case SampleConfigTypeNeman:
        {
            vector<double> layerVariances;
            for (auto layer : layerInfos_) {
                layerVariances.push_back(layerVariance(layer.second));
            }
            neymanAllocation(layerPopulationSizes, layerVariances, totalSampleSize, layerSampleSizes);
        }
        break;
        case SampleConfigTypeMetricScale:
            // layerMetricValues 包含每个层指标值的列表，如何计算？
            // indicatorProportionalAllocation(layerMetricValues, layerPopulationSizes, totalSampleSize, layerSampleSizes);
            return;
            break;
        case SampleConfigTypeSampleFrameScale:
            layerSampleSizes = proportionalAllocation(layerPopulationSizes, totalSampleSize);
            break;
        default:
            return;
        }

        assert(layerInfos_.size() == layerSampleSizes.size());

        unsigned i = 0;
        for(auto& layer : layerInfos_) {
            layer.second.sampleSize = layerSampleSizes[i++];
        }
    }

    

    /// 根据要求的最大误差计算每层的样本量, 并保存到layerSampleSizes_中
    bool StratifiedSample::calcSampleSize()
    {
        if (relativeError_ == 0)
            return false;

        sampleSize_ = 0;
        for (auto& layer : layerInfos_)
        {
            layer.second.sampleSize = std::min(layerPopulationSize(layer.second), calculateSampleSize(relativeError_, 0.95, layerStdDev(layer.second)));
            sampleSize_ += layer.second.sampleSize;
        }

        return true;
    }

    bool StratifiedSample::calcRelativeError()
    {
        if (sampleSize_ == 0) {
            return false;
        }

        for (auto& layer : layerInfos_)
        {
            layer.second.error = calculateMaxRelativeError(layer.second.sampleSize, 0.95, layerStdDev(layer.second));
        }
                
        return true;
    }
	
    unsigned StratifiedSample::layerPopulationSize(const LayerInfo& layer) const
    {
        return isStringLayering() ? layer.populationNos.size() : layer.populationSize;
    }

    double StratifiedSample::layerStdDev(const LayerInfo& layer) const
    {
        return sqrt(layerVariance(layer));
    }
    double StratifiedSample::layerVariance(const LayerInfo& layer) const
    {
        if (isStringLayering()) {
            return 0;
        }

        auto c = sampleFrame_.getColumn(layerMetric_);
        if (c == nullptr) {
            return 0;
        }

        auto metricType = sampleFrame_.getColumnType(layerMetric_); 
        vector<double> data;
        for(auto p : layer.populationNos) {
            if (metricType == Metric::TypeInt) {
                data.push_back(std::get<int>((*c)[p]));
            }
            else if (metricType == Metric::TypeDouble) {
                data.push_back(std::get<double>((*c)[p]));
            }
        }
        
         // 计算平均值
        double mean = std::accumulate(data.begin(), data.end(), 0.0) / data.size();
        
        // 计算方差
        return std::accumulate(data.begin(), data.end(), 0.0,
                                          [&](double acc, double val) { return acc + pow(val - mean, 2); }) / data.size();
    }
    

    SampleResult* StratifiedSample::sample() {
         if (sampleSize_ == 0 && 0 == relativeError_) {
            return nullptr;
        }
        if (metric_.empty()) {
            return nullptr;
        }

        SampleResult sampleResult;
        sampleResult.sampleDesign = this;
            
        // 确定每层样本情况
        for(auto layer : layerInfos_) {
            // 随机生成n个不重复的整数
            auto lps = layerPopulationSize(layer.second);
            sampleResult.results[layer.first] = generateUniqueRandomNumbers(lps, layer.second.sampleSize);
        }
        
        return &SampleManager::getInstance().addSampleResult(sampleResult);
    }    
}